import base64
from flask import request, Blueprint, render_template

from models.host import Host

from service.ssh import SSH
from service.authtoken import login_required

webssh_route = Blueprint('webssh', __name__)


@webssh_route.route('/webshell', methods=["GET"])
@login_required
def web_shell():
    host_id = request.args.get("host_id")
    if host_id:
        if Host.query.filter(Host.id == host_id).first():
            return render_template('webshell.html', host_id=host_id)
    return "主机不存在！"


@webssh_route.route('/webssh', methods=["GET"], websocket=True)
@login_required
def web_ssh():
    user_socket = request.environ.get("wsgi.websocket")
    try:
        if user_socket:
            host_id = request.args.get("host_id")
            pty_width = request.args.get("pty_width", 128)
            pty_height = request.args.get("pty_height", 64)

            host_obj = Host.query.filter(Host.id == host_id).first()
            if host_obj:
                ssh_obj = SSH(socket=user_socket)
                if host_obj.auth_type == "password":
                    password = base64.b64decode(host_obj.password).decode('utf-8')
                    ssh_obj.connect(host=host_obj.ip, username=host_obj.username, password=password, port=host_obj.port,
                                    pty_width=int(pty_width), pty_height=int(pty_height))
                else:
                    password = base64.b64decode(host_obj.password).decode('utf-8')
                    ssh_obj.connect(host=host_obj.ip, username=host_obj.username, password=password, port=host_obj.port,
                                    pty_width=int(pty_width), pty_height=int(pty_height))

                ssh_obj.socket_to_shell()
    except Exception as e:
        print(e)
        return 'Finished'
